# This Maximum-Entropy method, like any MEM method, has 2 regularizers:
# one for total flux and one for entropy. In an ideal situation, total 
# flux is fixed from the data (not here as we are in the high-contrast
# limit with an unresolved point source), and the MEM regularizer can
# be set to give a reduced chi-squared = 1. Often in a MEM technique, 
# the total flux is constrained using a prior. This is the approach
# that we will take here.

import numpy as np
import matplotlib.pyplot as plt
import pyfits

#Input a fits file, containing an linear image-plane (i.e. high-contrast)
#representation of a set of kernel phases, plus data and errors in 
#the fits extension
def data_in(filename):
    hdulist = pyfits.open(filename)
    pxscale = hdulist[0].header['PXSCALE']
    Amat = hdulist[0].data
    dimx = Amat.shape[2]
    dimy = Amat.shape[1]
    Amat = Amat.reshape([Amat.shape[0],dimx*dimy])
    tbdata = hdulist[1].data
    theta = tbdata[0]
    theta_sig = tbdata[1]
    return Amat,theta,theta_sig,dimx,dimy,pxscale

#Return the entropy functional.
def ent(vector,prior):
    return np.sum(vector/prior*np.log(vector/prior))

#Return the gradient of the entropy functional.
def grad_ent(vector,prior):
    return (np.log(vector/prior) + 1)/prior 

#Return the function we are trying to minimise.
def f(z,Amat,theta,theta_sig,alpha,prior):
    theta_mod = np.dot(Amat,z)
    chi2 = np.sum((theta_mod - theta)**2/theta_sig**2)
    stat = chi2 + alpha*ent(z,prior) 
    return stat,theta_mod,chi2

#Return the gradient of the function we are trying to minimise.
def grad_f(z,Amat,theta,theta_sig,theta_mod,alpha,prior):
    the_vect = (theta_mod - theta)/theta_sig**2
    return 2*np.dot(the_vect.reshape(1,theta.size),Amat)+alpha*grad_ent(z,prior)

#Return the function we are trying to minimise along a line 
#defined by z = s + p*t.
#Inputs include the matrix products A * s and A * p.
def fline(t,s,p,As,Ap,theta,theta_sig,alpha,prior):
    z = s + p*t
    chi2 = np.sum((As + Ap*t - theta)**2/theta_sig**2) + alpha*ent(z,prior)
    return chi2

#Perform a line search from a starting image "s" in a direction "p"
def line_search(s,p,As,Ap,theta,theta_sig,alpha,prior):
    #This is a golden section search with a fixed 30 iterations.
    #Returns the t value of the minimum.
    niter=30
    if (min(p) > 0):
        #This shouldn't really happen
        tmax = 2*np.max(s)/np.max(p)
    else:
        ts = -s/p.clip(-1e12,-1e-12)
        #A somewhat arbitrary limit below to prevent log(0)
        tmax = np.min(ts) * (1-1e-6)
    phi = (1 + np.sqrt(5))/2.0
    resphi = 2 - phi
    t1 = 0
    t3 = tmax
    t2 = t3/phi
    f1 = fline(t1,s,p,As,Ap,theta,theta_sig,alpha,prior)
    f2 = fline(t2,s,p,As,Ap,theta,theta_sig,alpha,prior)
    f3 = fline(t3,s,p,As,Ap,theta,theta_sig,alpha,prior)
    for i in range(niter):
        if (t3 - t2 > t2 - t1):
          t4 = t2 + resphi * (t3 - t2)
          f4 = fline(t4,s,p,As,Ap,theta,theta_sig,alpha,prior)
          if (f4 < f2):
              f1 = f2
              f2 = f4
              t1 = t2
              t2 = t4
          else:
              f3 = f4
              t3 = t4
        else:
          t4 = t2 - resphi * (t2 - t1)
          f4 = fline(t4,s,p,As,Ap,theta,theta_sig,alpha,prior)            
          if (f2 < f4):
              f1 = f4
              t1 = t4
          else:
              f3 = f2
              f2 = f4
              t3 = t2
              t2 = t4
    return [t1,t2,t3][np.argmin([f1,f2,f3])]

#Create a Maximum-Entropy image based on an input fits
#file and a prior. Optional inputs are:
# alpha: A starting value for the MEM functional multiplier (default=1.0)
# gain: The servo gain for adjusting alpha to achieve chi^2=1
# niter: The number of iterations. 
def mem_image(infile,prior,**kwargs):
    keys = kwargs.keys()
    if ('alpha' in keys):
        alpha = kwargs['alpha']
    else:
        alpha = 1.0
    if ('gain' in keys):
        gain = kwargs['gain']
    else:
        #A difference of 1.0 in chi^2 gives a difference of 0.1 in log(alpha)
        gain = 0.1
    if ('niter' in keys):
        niter = kwargs['niter']
    else:
        niter = 400

    [Amat,theta,theta_sig,dimx,dimy,pxscale] = data_in(infile)
    z = prior.copy()
    start_i=5 #Should this be an option?

    for i in range(niter):
        if ( (i+1) % 10 == 0):
            print 'Done: ' + str(i+1) + ' of ' + str(niter) +' iterations. Chi^2: ' + str(chi2) + ' alpha: ' +str(alpha)
        #While this loop runs, we need to adjust alpha so that chi^2=1.
        #Lower values of alpha give lower chi^2. Not knowing how to do this,
        #lets use a servo-loop approach.
        if (i >= start_i):
           err = 1 - chi2/len(theta) #The difference in chi-squared
           alpha = alpha*np.exp( gain*err )
        [stat,theta_mod,chi2] = f(z,Amat,theta,theta_sig,alpha,prior)
        #print chi2/len(theta), alpha
        [grad] = grad_f(z,Amat,theta,theta_sig,theta_mod,alpha,prior)
        #The following is the Polak-Ribiere Nonlinear conjugant gradient method. 
        #(http://en.wikipedia.org/wiki/Nonlinear_conjugate_gradient_method)
        #It seems to be only about a factor of 2 faster than steepest descent. 
        if (i == 0):
            dir = -grad.copy() #The direction of the minimum
            p = dir.copy()
        else:
            lastdir = dir.copy()
            dir = -grad.copy() #The direction of the minimum  
            beta = np.sum(dir*(dir - lastdir))/np.sum(dir*lastdir)
            beta = np.max(beta,0)
            p = dir + beta*p
        s = z.copy()
        As = theta_mod 
        Ap = np.dot(Amat,p)
        tmin = line_search(s,p,As,Ap,theta,theta_sig,alpha,prior)
        z = s + tmin*p
    [stat,theta_mod,chi2] = f(z,Amat,theta,theta_sig,0,prior)
    print "Chi2: " + str(chi2/len(theta))
    extent = [dimx/2*pxscale,-dimx/2*pxscale,-dimy/2*pxscale,dimy/2*pxscale]
    return [z,extent]

#Here is some example code to create an image of LkCa15 in K-band.
prior = np.ones(81**2.0)/81**2.0/50
[z,extent] = mem_image('LkCa15_K2010_implane.fits',prior)
im = z.reshape(81,81)
plt.imshow(im, interpolation='nearest',cmap=plt.get_cmap('gist_heat'),extent=extent)
plt.plot(0,0,'w*', ms=15)
plt.axis(extent)
plt.xlabel('Delta RA (milli-arcsec)', fontsize='large')
plt.ylabel('Delta Dec (milli-arcsec)', fontsize='large')
print "Total contrast (mags): " + str(-2.5*np.log10(np.sum(z)))




